from flask import Flask, request, jsonify, render_template
import numpy as np
from ultralytics import YOLO
import os
from werkzeug.utils import secure_filename
from flask_cors import CORS  # 导入CORS
import logging
import tempfile

# 设置日志级别为DEBUG
logging.basicConfig(level=logging.DEBUG)

# 创建 Flask 应用
app = Flask(__name__, template_folder='templates', static_url_path='', static_folder='')  # 空字符串表示根目录
CORS(app)  # 启用CORS

# 加载 YOLO 模型
model = YOLO("recognize/best.pt")  # 确保路径正确

# 允许的文件类型
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/')
def index():
    """渲染主页"""
    return render_template('recognize.html')

@app.route('/api/recognize', methods=['POST'])
def predict():
    """处理文件上传和预测"""
    if 'image' not in request.files:
        logging.error('No file part in the request')
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['image']
    
    if file.filename == '':
        logging.error('No selected file')
        return jsonify({'error': 'No selected file'}), 400
    
    # 打印文件信息用于调试
    logging.debug(f"File: {file}")
    logging.debug(f"Filename: {file.filename}")
    logging.debug(f"File type: {file.content_type}")
    
    if file and allowed_file(file.filename):
        # 使用 tempfile 模块生成临时文件
        with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
            filepath = temp_file.name
            file.save(filepath)
        
        # 使用 YOLO 模型进行预测
        results = model(filepath)
        names_dict = results[0].names
        probs = results[0].probs.data.tolist()
        
        # 打印预测结果用于调试
        logging.debug(f"Prediction names: {names_dict}")
        logging.debug(f"Probabilities: {probs}")

        if not probs or max(probs) <= 0:
            logging.error('No valid prediction found')
            os.remove(filepath)
            return jsonify({'error': 'No valid prediction found'}), 400
        
        # 构建返回结果
        result = {
            'prediction': names_dict[np.argmax(probs)],
            'probability': max(probs),
        }
        
        # 删除临时文件
        os.remove(filepath)
        return jsonify(result)
    else:
        logging.error('File type not allowed')
        return jsonify({'error': 'File type not allowed'}), 400

# 如果直接运行此脚本，则启动Flask开发服务器
if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)  # 运行在5000端口